[MoE] Use the sorted fused-MoE path for host-resident expert weights - #62054

Overview

At small batch (num_tokens * top_k * 4 <= num_experts, e.g. up to 4 tokens for Qwen3 MoE) the Triton fused-MoE path skips moe_align_block_size and gives every (token, expert) pair its own block (#29354, #39016). That saves the alignment kernel when weights are in device memory. With UVA weight offloading (--offload-backend uva --cpu-offload-params experts) the expert weights are read from host memory in place, and an expert routed to several tokens is then read over PCIe once per token. This PR takes the sorted path when the expert weights carry the UVA offload marker, so each routed expert is read once per step. Device-resident weights keep the naive path.

Claims

On RTX 4090 (PCIe 4.0 x16), natural-text chat requests, 128 output tokens:

• Qwen3-30B-A3B-Instruct-2507-FP8, one GPU, 16 GiB of experts offloaded: +7% / +10% / +15–19% output throughput at 2 / 3 / 4 requests in flight; unchanged at 1, 8 and 16.
• Qwen3-235B-A22B-Instruct-2507-FP8, TP8, 16 GiB per GPU offloaded: +7% / +10% / +18% at 2 / 3 / 4; unchanged at 1.
• Base is the vLLM 0.31.0 wheel, whose two touched files are byte-identical to the base commit. Measured on CUDA only; ROCm's Triton path takes the same branch.
• No change without offloading: the code path is the same; measured on the 30B at TP2 with weights in VRAM (one start per build).
• Quality: GSM8K (lm-eval 0.4.13, 5-shot, greedy, first 250 questions, 4 requests in flight so decode steps take the changed path), 30B with 16 GiB offloaded: 0.888 flexible / 0.880 strict exact match for both builds, and all 250 responses byte-identical.
• Output: greedy completions identical to unpatched on the 30B (all starts). On the 235B at TP8 greedy output varies from start to start without the patch (three unpatched starts gave three different completions and three patched starts two, with one completion common to both builds), so it cannot distinguish the builds there.

Validation

• pytest tests/kernels/moe/test_moe.py -k "host_resident_weights or naive_block_assignment": 44 passed on an RTX 4090, on the original base and again after rebasing onto current main (nightly wheel of the base commit, patched files applied). The new test_host_resident_weights_take_sorted_path runs UVA-mapped weights (and device weights as the control) at 1/2/4 tokens through both the modular TritonExperts path and the legacy fused_moe, asserts that moe_align_block_size is called exactly when the weights are host-resident, and compares the output with torch_moe (inputs scaled to O(1) outputs, atol=rtol=2e-2). Without the change its 6 host-resident cases fail.

• Every applicable pre-commit hook passes: ruff check/format, mypy 3.10–3.13, check_torch_cuda, test tethering, typos, SPDX, lazy imports, forbidden imports, boolean context manager.

• End to end, base vLLM against this PR, one vllm serve, requests from a client sending natural-language chat prompts (32 topics × 8 tasks, temperature 0.7, a fixed seed per request, ignore_eos):

  setup	in flight	starts (base / PR)	output tok/s, base	output tok/s, this PR	change
  30B, 1 GPU, 16 GiB on host	1	2 / 2	20.7	20.5–20.7	−0.5% (within spread)
  	2	2 / 2	22.2	23.8	+7%
  	3	2 / 2	22.6	24.9–25.0	+10%
  	4	2 / 2	23.1	26.5–27.4	+15–19%
  	8	2 / 2	34.5–34.9	33.5–35.0	within spread
  	16	2 / 2	49.0–49.1	48.5–49.1	within spread
  235B, TP8, 16 GiB/GPU on host	1	1 / 1	16.6	16.6	0%
  	2	3 / 3	17.9–18.3	19.1–19.5	+7%
  	3	3 / 3	19.0	20.9	+10%
  	4	3 / 3	19.7	23.2–23.3	+18%
  30B, TP2, all in VRAM	1–4	1 / 1	196.8 / 300.3 / 364.1 / 474.8	197.0 / 300.6 / 363.0 / 473.5	within 0.3%
  	8	1 / 1	741.5	727.1	−1.9% (one start; both builds run the same code here)

• Reproducible with vLLM's own benchmark (vllm bench serve --dataset-name sharegpt --sharegpt-output-len 128 --ignore-eos --num-prompts N --max-concurrency c --seed 1, ShareGPT_V3_unfiltered_cleaned_split; N = 8 for c < 4, else 2c), 30B on one GPU with 16 GiB offloaded, two starts each:

  --max-concurrency	output tok/s, base	output tok/s, this PR	change (midpoints)
  1	20.3	20.3	0%
  2	21.3–22.2	23.4–23.6	+8%
  3	21.6–22.6	25.3–25.4	+15%
  4	22.2–23.0	25.6–26.7	+16%

# 30B, one GPU (GPU 0 on NUMA node 3 here), 16 GiB of experts on the host
CUDA_VISIBLE_DEVICES=0 numactl --cpunodebind=3 --membind=3 vllm serve Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \
  --max-model-len 4096 --gpu-memory-utilization 0.90 --offload-backend uva --cpu-offload-gb 16 --cpu-offload-params experts
# 235B, TP8
vllm serve Qwen/Qwen3-235B-A22B-Instruct-2507-FP8 --tensor-parallel-size 8 --max-model-len 4096 \
  --gpu-memory-utilization 0.92 --offload-backend uva --cpu-offload-gb 16 --cpu-offload-params experts --kv-cache-dtype fp8

Details

Why. With experts in host memory, a decode step's cost is dominated by the bytes it reads over PCIe, and on the naive path an expert routed to several of the step's tokens is read once per token. The sorted path reads each routed expert once per step. How much that saves depends on how often a step's tokens share experts: on natural-text prompts the gain is in the Validation table; on random-token prompts with greedy decoding, whose routing is far more repetitive, the same change doubles throughput at 4 requests in flight. We report natural text.

Implementation.

• _prepare_expert_assignment takes weights_host_resident (default False), which disables the naive block assignment.
• TritonExperts.apply and the legacy fused_experts_impl pass _weights_host_resident(w1, w2): true when either weight has _vllm_is_uva_offloaded, which the UVA offloader sets (and device_loading_context re-sets after process_weights_after_loading). w1/w2 are the layer's own parameters on this path, so it is a getattr per forward inside the moe_forward custom op; the result depends only on the parameters, so CUDA-graph capture per size is unaffected. "Either", because UVA offloads per parameter and a budget cut-off can leave w13 on the host and w2 on the device.
• Other MoE paths are unaffected: Marlin and TritonWNA16Experts always align; the prefetch offloader copies weights to the device before compute and never sets the marker. Two model-specific callers of _prepare_expert_assignment (models/k2_horizon.py, models/dots3_note/nvidia/vision_moe.py) do not pass the flag and keep today's behaviour; they can opt in separately.
• Why a per-forward check rather than a decision when the kernel is built: replace_parameter in process_weights_after_loading installs a new Parameter without the marker, and device_loading_context re-marks it only afterwards, so a setup-time check would read False.

Related. No open or merged PR makes this change (searched 2026-10-07, rechecked 2026-10-16). The naive path came from #29354 and #39016; #34920 (closed, stale) proposed making it tunable from configs; #55268 (open) warms up Triton MoE variants per batch size, and this PR changes which variant host-resident weights use at ≤ 4 tokens. #57215 (open) keeps MoE experts offloaded during weight post-processing on XPU, the marker lifecycle this PR relies on; #56177 (open) is a shared GPU expert pool. RFC #38256 discusses expert caches for offloaded MoE, where a comment reports the same effect with a GPU-side slot buffer (the naive path stops sharing an expert's weight tile between tokens at batch ≥ 2). #37190, #57794 and #57943 decide which experts stay in VRAM; this PR is independent of that and helps whichever experts remain on the host.

AI assistance. Claude Opus 5.5 assisted with the code, tests, benchmarks and this description; the tests and benchmarks above were run on my hardware. I reviewed every changed line and the results, and can answer for it; the commit carries an Assisted-by: trailer.